package com.lm.tokenbucket.example;
import java.util.concurrent.atomic.AtomicLong;

public class TokenBucketExample {

    private final long capacity; // 桶的容量
    private final long refillRate; // 令牌生成速率，单位：令牌/秒
    private AtomicLong tokens; // 当前令牌数量
    private long lastRefillTimestamp; // 上次填充令牌的时间戳

    public TokenBucketExample(long capacity, long refillRate) {
        this.capacity = capacity;
        this.refillRate = refillRate;
        this.tokens = new AtomicLong(capacity);
        this.lastRefillTimestamp = System.nanoTime();
    }

    private synchronized void refill() {
        long now = System.nanoTime();
        long elapsedTime = now - lastRefillTimestamp;
        long tokensToAdd = (elapsedTime * refillRate) / 1_000_000_000L;
        if (tokensToAdd > 0) {
            tokens.set(Math.min(capacity, tokens.get() + tokensToAdd));
            lastRefillTimestamp = now;
        }
    }

    public synchronized boolean tryConsume(long numTokens) {
        refill();
        if (tokens.get() >= numTokens) {
            tokens.addAndGet(-numTokens);
            return true;
        }
        return false;
    }

    public static void main(String[] args) {
        TokenBucketExample tokenBucket = new TokenBucketExample(10, 5); // 容量为10，速率为5令牌/秒

        for (int i = 0; i < 20; i++) {
            if (tokenBucket.tryConsume(1)) {
                System.out.println("请求成功通过: " + i);
            } else {
                System.out.println("请求被拒绝: " + i);
            }

            try {
                Thread.sleep(200); // 模拟每200ms发送一个请求
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }
    }
}
